// @HEADER
//*********************************************************************//
//  SiYuan: A numerical PDE solver                                     //
//  Copyright (2022) YUAN Xi                                           //
//  This Software is released under the BSD 2-Clause license detailed  //
//  in the file "LICENSE" in the top-level SiYuan directory            //
//*********************************************************************//
// @HEADER

#ifndef DIFFUSION_EQUATION_IMPL_HPP
#define DIFFUSION_EQUATION_IMPL_HPP

#include "Teuchos_ParameterList.hpp"
#include "Teuchos_StandardParameterEntryValidators.hpp"
#include "Teuchos_RCP.hpp"
#include "Teuchos_Assert.hpp"
#include "Phalanx_DataLayout_MDALayout.hpp"
#include "Phalanx_FieldManager.hpp"

#include "Panzer_IntegrationRule.hpp"
#include "Panzer_BasisIRLayout.hpp"

// include evaluators here
#include "Panzer_Integrator_BasisTimesScalar.hpp"
#include "Panzer_Integrator_TransientBasisTimesScalar.hpp"
#include "Panzer_Integrator_GradBasisDotVector.hpp"
#include "Panzer_ScalarToVector.hpp"
#include "Panzer_Sum.hpp"
#include "Panzer_Constant.hpp"
#include "EquationSet_Impl.hpp"
#include "Convection.hpp"

// ***********************************************************************
template <typename EvalT>
SiYuan::EquationSet_Diffusion<EvalT>::
EquationSet_Diffusion(const Teuchos::RCP<Teuchos::ParameterList>& params,
                   const int& default_integration_order,
                   const panzer::CellData& cell_data,
                   const Teuchos::RCP<panzer::GlobalData>& global_data,
                   const bool build_transient_support) :
  EquationSet<EvalT>(params,default_integration_order,cell_data,global_data,build_transient_support )
{
  // ********************
  // Validate and parse parameter list
  // ********************
  {    
    Teuchos::ParameterList valid_parameters;
    this->setDefaultValidParameters(valid_parameters);

    valid_parameters.set("Model ID","","Closure model id associated with this equaiton set");
    valid_parameters.set("Prefix","","Prefix for using multiple instatiations of thei equation set");
    valid_parameters.set("DOF Name","Temperature","DOF Name");
    valid_parameters.set("Material Name","?","Material Name");
    valid_parameters.set("Basis Type","HGrad","Type of Basis to use");
    valid_parameters.set("Basis Order",1,"Order of the basis");
    valid_parameters.set("Integration Order",-1,"Order of the integration rule");
    valid_parameters.set<bool>("Do convection",false,"Enables or disables convection term");
    valid_parameters.set<bool>("Has Source",false,"If have source terms");

    params->validateParametersAndSetDefaults(valid_parameters);
  }

  do_convection = params->get<bool>("Do convection");
  m_prefix = params->get<std::string>("Prefix");
  m_dof_name = params->get<std::string>("DOF Name");
  m_matl_name = params->get<std::string>("Material Name");
  std::string basis_type = params->get<std::string>("Basis Type");
  int basis_order = params->get<int>("Basis Order");
  std::string model_id = params->get<std::string>("Model ID");
  int integration_order = params->get<int>("Integration Order");

  // ********************
  // Setup DOFs and closure models
  // ********************
  {
    m_dof_name = m_prefix+m_dof_name;

    this->addDOF(m_dof_name,basis_type,basis_order,integration_order);
    this->addDOFGrad(m_dof_name);
    if (this->buildTransientSupport()) {
      this->addDOFTimeDerivative(m_dof_name);
      this->addC2("Density");
      this->addC2("Specific Heat");
    }
  }

  this->addClosureModel(model_id);
  this->setupDOFs();
  has_source = params->get<bool>("Has Source");

  this->addC0("Thermal Conductivity");
}

// ***********************************************************************
template <typename EvalT>
void SiYuan::EquationSet_Diffusion<EvalT>::
buildAndRegisterEquationSetEvaluators(PHX::FieldManager<panzer::Traits>& fm,
                                      const panzer::FieldLibrary& /* fl */,
                                      const Teuchos::ParameterList& /* user_data */) const
{
  using panzer::EvaluatorStyle;

  Teuchos::RCP<panzer::IntegrationRule> ir = this->getIntRuleForDOF(m_dof_name); 
  Teuchos::RCP<panzer::BasisIRLayout> basis = this->getBasisIRLayoutForDOF(m_dof_name); 

  // Transient Operator
  if (this->buildTransientSupport())
  {
    std::string resName("RESIDUAL_" + m_dof_name),
           valName("DXDT_" + m_dof_name);
    double multiplier(1);
  //  std::vector<std::string> fieldMultipliers{m_prefix + "Density",
  //    m_prefix + "Specific Heat"};
    Teuchos::RCP<PHX::Evaluator<panzer::Traits>> op =Teuchos:: rcp(new
      panzer::Integrator_BasisTimesScalar<EvalT, panzer::Traits>(EvaluatorStyle::CONTRIBUTES,
      resName, valName, *basis, *ir, multiplier, this->getC2()) );
    this->template registerEvaluator<EvalT>(fm, op);
  }

  // Diffusion Operator
  {
    Teuchos::ParameterList p("Diffusion Residual");
    p.set("Residual Name", "RESIDUAL_"+m_dof_name);
    p.set("Flux Name", "GRAD_"+m_dof_name);
    p.set("Basis", basis);
    p.set("IR", ir);
    Teuchos::RCP<const std::vector<std::string> > vec = 
      Teuchos::rcp(new std::vector<std::string>(this->getC0()));
    p.set("Field Multipliers", vec);
    p.set("Multiplier", 1.0);
    
    Teuchos::RCP< PHX::Evaluator<panzer::Traits> > op = 
      Teuchos::rcp(new panzer::Integrator_GradBasisDotVector<EvalT,panzer::Traits>(p));

    this->template registerEvaluator<EvalT>(fm, op);
  }
  
  // Convection Operator
  if (do_convection) {

    // Combine scalar velocities into a velocity vector
    {
      Teuchos::ParameterList p("Velocity: ScalarToVector");
      Teuchos::RCP<std::vector<std::string> > scalar_names = Teuchos::rcp(new std::vector<std::string>);
      scalar_names->push_back(m_prefix+"UX");
      scalar_names->push_back(m_prefix+"UY");
      p.set<Teuchos::RCP<const std::vector<std::string> > >("Scalar Names", scalar_names);
      p.set("Vector Name", m_prefix+"U");
      p.set("Data Layout Scalar",ir->dl_scalar);
      p.set("Data Layout Vector",ir->dl_vector);

      Teuchos::RCP< PHX::Evaluator<panzer::Traits> > op = 
        Teuchos::rcp(new panzer::ScalarToVector<EvalT,panzer::Traits>(p));
      
      this->template registerEvaluator<EvalT>(fm, op);
    }

    // Evaluator to assemble convection term
    {
      Teuchos::ParameterList p("Convection Operator");
      p.set("IR", ir);
      p.set("Operator Name", m_dof_name+"_CONVECTION_OP");
      p.set("A Name", m_prefix+"U");
      p.set("Gradient Name", "GRAD_"+m_dof_name);
      p.set("Multiplier", 1.0);

      Teuchos::RCP< PHX::Evaluator<panzer::Traits> > op = 
        Teuchos::rcp(new Convection<EvalT,panzer::Traits>(p));
      
      this->template registerEvaluator<EvalT>(fm, op);
    }

    // Integration operator (could sum this into source for efficiency)
    {
      std::string resName("RESIDUAL_" + m_dof_name),
             valName(m_dof_name + "_CONVECTION_OP");
      double multiplier(1);
    //  std::vector<std::string> fieldMultipliers{m_prefix + "Density",
    //    m_prefix + "Specific Heat"};
      Teuchos::RCP<PHX::Evaluator<panzer::Traits>> op = Teuchos::rcp(new
        panzer::Integrator_BasisTimesScalar<EvalT, panzer::Traits>(EvaluatorStyle::CONTRIBUTES,
        resName, valName, *basis, *ir, multiplier, this->getC2()));
      this->template registerEvaluator<EvalT>(fm, op);
    }
  }

  // Source Operator
  if( has_source ) {   
    std::string resName("RESIDUAL_" + m_dof_name),
           valName("SOURCE_" + m_dof_name);
    double multiplier(-1);
    Teuchos::RCP<PHX::Evaluator<panzer::Traits>> op = Teuchos::rcp(new
      panzer::Integrator_BasisTimesScalar<EvalT, panzer::Traits>(EvaluatorStyle::CONTRIBUTES,
      resName, valName, *basis, *ir, multiplier));
    this->template registerEvaluator<EvalT>(fm, op);
  }
}

// ***********************************************************************

#endif
